import torch
import torch.nn as nn
import numpy as np
import torchvision
import math
class CustomLSTM(nn.Module):
    def __init__(self, input_sz, hidden_sz):
        super().__init__()
        self.input_sz = input_sz
        self.hidden_size = hidden_sz
        self.W = nn.Parameter(torch.Tensor(input_sz, hidden_sz * 4))
        self.U = nn.Parameter(torch.Tensor(hidden_sz, hidden_sz * 4))
        self.bias = nn.Parameter(torch.Tensor(hidden_sz * 4))
        self.init_weights()
 
    def init_weights(self):
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)
 
    def forward(self, x, 
                init_states=None):
        """Assumes x is of shape (batch, sequence, feature)"""
        bs, seq_sz, _ = x.size()
        hidden_seq = []
        if init_states is None:
            h_t, c_t = (torch.zeros(bs, self.hidden_size).to(x.device), 
                        torch.zeros(bs, self.hidden_size).to(x.device))
        else:
            h_t, c_t = init_states
 
        HS = self.hidden_size
        for t in range(seq_sz):
            x_t = x[:, t, :]
            # batch the computations into a single matrix multiplication
            gates = x_t @ self.W + h_t @ self.U + self.bias
            i_t, f_t, g_t, o_t = (
                torch.sigmoid(gates[:, :HS]), # input
                torch.sigmoid(gates[:, HS:HS*2]), # forget
                torch.tanh(gates[:, HS*2:HS*3]),
                torch.sigmoid(gates[:, HS*3:]), # output
            )
            c_t = f_t * c_t + i_t * g_t
            h_t = o_t * torch.tanh(c_t)
            hidden_seq.append(h_t.unsqueeze(0))
        hidden_seq = torch.cat(hidden_seq, dim=0)
        # reshape from shape (sequence, batch, feature) to (batch, sequence, feature)
        hidden_seq = hidden_seq.transpose(0, 1).contiguous()
        return hidden_seq, (h_t, c_t)

class Anti_spoof_net_RNN(nn.Module):
  
    def __init__ (self,device):
        super(Anti_spoof_net_RNN,self).__init__()
        
        self.hidden_dim=100
        self.input_dim=32*32
        self.num_layers=1
        self.batch_size=1

        # self.hidden = (torch.zeros(self.batch_size, self.hidden_dim).to(device), torch.zeros( self.batch_size, self.hidden_dim).to(device))
        # self.hidden = (torch.zeros(self.num_layers, self.batch_size, self.hidden_dim), torch.zeros(self.num_layers, self.batch_size, self.hidden_dim))
        self.LSTM=nn.LSTM(input_size=self.input_dim,hidden_size=self.hidden_dim,num_layers=self.num_layers)
        self.LSTM = CustomLSTM(self.input_dim,self.hidden_dim)
        self.fc=nn.Linear(self.hidden_dim,2)
        
    def forward(self,F):
            
        #F est de dimension [5,32,32,1]
        F = F.view(F.shape[0],1,-1)
        lstm_out, _ = self.LSTM(F) #lstm_out[5,1,100]
        
        F = self.fc(lstm_out) #F[5,1,2]
        # R = torch.fft.fft(R, signal_ndim=1, normalized=False)
        # F = torch.fft.fft(F, dim=1, norm='ortho')
        return F #F[5,1,2]
